#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简易RAG系统 - 教学演示版本
目标：30分钟内看到RAG效果，理解核心概念

设计理念：
- 极简实现，先看效果
- 代码集中在一个文件
- 注释详细，标注优化方向
"""

import os
import re
import json
import numpy as np
from typing import List, Dict, Tuple, Any
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import jieba
import sys

# 添加项目路径
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)

class SimpleRAGSystem:
    """
    简易RAG系统
    
    功能：文档加载 → 切分 → 向量化 → 存储 → 检索 → 生成回答
    
    优化方向：
    1. 文档处理：支持PDF/Word → 使用专业解析库
    2. 文档切分：智能切分 → 基于语义的切分
    3. 向量化：高质量向量 → 使用BERT/Sentence-BERT
    4. 存储：持久化存储 → 使用向量数据库
    5. 检索：多路检索 → 混合检索+重排序
    6. 生成：优化提示词 → 使用专业的Prompt工程
    """
    
    def __init__(self):
        """初始化RAG系统"""
        print("🚀 初始化简易RAG系统")
        
        # 存储文档和向量
        self.documents = []  # 存储文档片段
        self.document_vectors = []  # 存储文档向量
        self.vectorizer = None  # TF-IDF向量化器
        
        # 配置参数
        self.chunk_size = 300  # 文档切分大小（字符数）
        self.chunk_overlap = 50  # 重叠大小
        self.top_k = 3  # 检索返回的文档数量
        
        print("✅ RAG系统初始化完成")
    
    def load_documents_from_folder(self, folder_path: str) -> List[str]:
        """
        从文件夹加载文档（当前版本：仅支持.txt文件）
        
        优化方向：
        - 支持PDF解析：使用PyPDF2、pdfplumber
        - 支持Word解析：使用python-docx
        - 支持更多格式：Excel、PowerPoint等
        """
        print(f"📁 从文件夹加载文档: {folder_path}")
        
        documents = []
        supported_extensions = ['.txt', '.md', '.py']
        
        if not os.path.exists(folder_path):
            print(f"❌ 文件夹不存在: {folder_path}")
            return documents
        
        for filename in os.listdir(folder_path):
            file_path = os.path.join(folder_path, filename)
            
            # 检查文件扩展名
            _, ext = os.path.splitext(filename)
            if ext.lower() not in supported_extensions:
                continue
            
            try:
                print(f"  📄 读取文件: {filename}")
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    if content.strip():
                        documents.append({
                            'filename': filename,
                            'content': content,
                            'source': file_path
                        })
                        print(f"    ✅ 成功读取 {len(content)} 个字符")
            
            except Exception as e:
                print(f"    ❌ 读取失败: {str(e)}")
        
        print(f"📊 总共加载了 {len(documents)} 个文档")
        return documents
    
    def simple_text_clean(self, text: str) -> str:
        """
        简单文本清洗
        
        优化方向：
        - 深度清洗：去除特殊字符、标点符号处理
        - 格式化处理：表格、列表结构保持
        - 语言检测：中英文混合处理
        """
        # 去除多余空白
        text = re.sub(r'\s+', ' ', text)
        text = text.strip()
        
        # 去除特殊字符（保留基本标点）
        text = re.sub(r'[^\w\s\u4e00-\u9fff，。！？；：""''（）【】_-]', '', text)
        
        return text
    
    def simple_chunk_split(self, text: str) -> List[str]:
        """
        简单文档切分（固定长度切分）
        
        优化方向：
        - 智能切分：按句子、段落切分
        - 语义切分：基于语义相似度切分
        - 结构化切分：保持文档结构（标题、列表等）
        """
        chunks = []
        text_length = len(text)
        
        for i in range(0, text_length, self.chunk_size - self.chunk_overlap):
            chunk = text[i:i + self.chunk_size]
            
            # 简单的边界处理：尽量在句号处切分
            if i + self.chunk_size < text_length:
                # 寻找最近的句号
                for j in range(len(chunk) - 1, max(0, len(chunk) - 100), -1):
                    if chunk[j] in '。！？\n':
                        chunk = chunk[:j + 1]
                        break
            
            if chunk.strip():
                chunks.append(chunk.strip())
        
        return chunks
    
    def build_vector_index(self, documents: List[Dict[str, str]]):
        """
        构建向量索引（当前版本：使用TF-IDF）
        
        优化方向：
        - 使用预训练模型：BERT、Sentence-BERT
        - 使用专业向量化API：OpenAI Embedding、阿里云百炼
        - 向量数据库：Chroma、Milvus、Faiss
        """
        print("🔢 开始构建向量索引...")
        
        # 1. 处理所有文档
        all_chunks = []
        chunk_metadata = []
        
        for doc in documents:
            print(f"  处理文档: {doc['filename']}")
            
            # 清洗文本
            cleaned_text = self.simple_text_clean(doc['content'])
            
            # 切分文档
            chunks = self.simple_chunk_split(cleaned_text)
            
            for i, chunk in enumerate(chunks):
                all_chunks.append(chunk)
                chunk_metadata.append({
                    'filename': doc['filename'],
                    'chunk_id': i,
                    'source': doc['source'],
                    'content': chunk
                })
        
        print(f"  📄 总共生成了 {len(all_chunks)} 个文档块")
        
        # 2. 向量化（使用TF-IDF）
        print("  🔢 开始向量化...")
        
        # 使用jieba分词（中文支持）
        def chinese_tokenizer(text):
            return list(jieba.cut(text))
        
        self.vectorizer = TfidfVectorizer(
            tokenizer=chinese_tokenizer,
            lowercase=False,
            max_features=5000,  # 最多5000个特征
            min_df=1,          # 最少出现1次
            max_df=0.95        # 最多出现在95%的文档中
        )
        
        # 训练并转换
        self.document_vectors = self.vectorizer.fit_transform(all_chunks)
        self.documents = chunk_metadata
        
        print(f"✅ 向量索引构建完成")
        print(f"  📊 向量维度: {self.document_vectors.shape}")
        print(f"  📄 文档块数量: {len(self.documents)}")
    
    def retrieve_relevant_docs(self, query: str, top_k: int = None) -> List[Dict[str, Any]]:
        """
        检索相关文档
        
        优化方向：
        - 多路检索：向量检索+关键词检索+语义检索
        - 重排序：使用重排序模型提高精度
        - 过滤机制：时间过滤、来源过滤等
        """
        if top_k is None:
            top_k = self.top_k
        
        if not self.vectorizer or len(self.documents) == 0:
            print("❌ 向量索引未构建，请先加载文档")
            return []
        
        print(f"🔍 检索相关文档: {query}")
        
        # 1. 查询向量化
        query_cleaned = self.simple_text_clean(query)
        query_vector = self.vectorizer.transform([query_cleaned])
        
        # 2. 计算相似度
        similarities = cosine_similarity(query_vector, self.document_vectors)[0]
        
        # 3. 获取最相似的文档
        top_indices = np.argsort(similarities)[::-1][:top_k]
        
        results = []
        for i, idx in enumerate(top_indices):
            doc = self.documents[idx].copy()
            doc['similarity'] = float(similarities[idx])
            doc['rank'] = i + 1
            results.append(doc)
        
        print(f"  📋 找到 {len(results)} 个相关文档")
        for result in results:
            print(f"    {result['rank']}. 相似度: {result['similarity']:.3f} - {result['filename']}")
        
        return results
    
    def generate_answer_with_llm(self, query: str, relevant_docs: List[Dict[str, Any]]) -> str:
        """
        使用LLM生成回答（集成阿里云百炼API）
        
        优化方向：
        - 更好的Prompt工程
        - 多轮对话支持
        - 回答质量评估
        """
        print("🤖 尝试使用LLM生成回答...")
        
        try:
            # 尝试导入项目中的LLM工具
            from utils.mul_tongyi import MultiTongyiClient
            
            client = MultiTongyiClient()
            
            # 构建上下文
            context_parts = []
            for doc in relevant_docs:
                context_parts.append(f"文档：{doc['filename']}\n内容：{doc['content'][:300]}...")
            
            context = "\n\n".join(context_parts)
            
            # 构建简洁的提示词
            prompt = f"""基于以下文档内容回答问题：

{context}

问题：{query}

请简洁准确地回答："""
            
            # 调用LLM
            response = client.call_qwen_api(prompt, "qwen-plus")
            
            if response and 'output' in response and 'text' in response['output']:
                answer = response['output']['text'].strip()
                
                # 添加引用信息
                answer += "\n\n📚 参考文档：\n"
                for doc in relevant_docs:
                    answer += f"- {doc['filename']} (相似度: {doc['similarity']:.3f})\n"
                
                print("✅ LLM回答生成成功")
                return answer
            else:
                print("⚠️ LLM响应格式异常，使用简单模板")
                return self.generate_simple_answer(query, relevant_docs)
        
        except Exception as e:
            print(f"⚠️ LLM调用失败，使用简单模板: {str(e)}")
            return self.generate_simple_answer(query, relevant_docs)
    
    def generate_simple_answer(self, query: str, relevant_docs: List[Dict[str, Any]]) -> str:
        """
        生成简单回答（模板拼接）
        """
        if not relevant_docs:
            return "抱歉，没有找到相关信息来回答您的问题。"
        
        # 构建上下文
        context_parts = []
        for doc in relevant_docs:
            context_parts.append(f"来源：{doc['filename']}\n内容：{doc['content'][:200]}...")
        
        context = "\n\n".join(context_parts)
        
        # 简单的模板回答
        answer = f"""基于检索到的文档内容，我找到了以下相关信息：

{context}

针对您的问题 "{query}"，根据上述文档内容进行回答。

相关文档：
"""
        
        for doc in relevant_docs:
            answer += f"- {doc['filename']} (相似度: {doc['similarity']:.3f})\n"
        
        return answer
    
    def chat(self, query: str, use_llm: bool = True) -> Dict[str, Any]:
        """
        完整的RAG对话流程
        
        Args:
            query: 用户问题
            use_llm: 是否使用LLM生成回答
        """
        print(f"\n" + "="*60)
        print(f"用户问题: {query}")
        print("="*60)
        
        # 1. 检索相关文档
        relevant_docs = self.retrieve_relevant_docs(query)
        
        # 2. 生成回答
        if use_llm:
            answer = self.generate_answer_with_llm(query, relevant_docs)
        else:
            answer = self.generate_simple_answer(query, relevant_docs)
        
        # 3. 返回结果
        result = {
            'query': query,
            'answer': answer,
            'relevant_docs': relevant_docs,
            'doc_count': len(relevant_docs)
        }
        
        print(f"\n🎯 RAG回答:")
        print(answer)
        print("="*60)
        
        return result
    
    def get_system_stats(self) -> Dict[str, Any]:
        """获取系统统计信息"""
        return {
            'total_documents': len(self.documents),
            'vector_dimension': self.document_vectors.shape[1] if self.document_vectors.size > 0 else 0,
            'chunk_size': self.chunk_size,
            'chunk_overlap': self.chunk_overlap,
            'top_k': self.top_k
        }

def demo_simple_rag():
    """演示简易RAG系统"""
    
    print("🎯 简易RAG系统演示")
    print("="*60)
    
    # 1. 创建RAG系统
    rag_system = SimpleRAGSystem()
    
    # 2. 加载文档（使用test_documents文件夹）
    documents_folder = os.path.join(os.path.dirname(__file__), 'test_documents', 'txt')
    documents = rag_system.load_documents_from_folder(documents_folder)
    
    if not documents:
        print("❌ 没有找到文档，请检查test_documents/txt文件夹")
        return
    
    # 3. 构建向量索引
    rag_system.build_vector_index(documents)
    
    # 4. 系统统计信息
    stats = rag_system.get_system_stats()
    print(f"\n📊 系统统计:")
    for key, value in stats.items():
        print(f"  {key}: {value}")
    
    # 5. 测试问答
    test_questions = [
        "什么是人工智能？",
        "机器学习的应用有哪些？",
        "Python编程有什么特点？",
        "员工数据包含哪些信息？",
        "如何进行招聘？"
    ]
    
    print(f"\n🧪 开始测试问答...")
    
    for question in test_questions:
        result = rag_system.chat(question, use_llm=False)  # 先使用简单模式演示
        
        # 简单的结果展示
        print(f"\n问题: {result['query']}")
        print(f"找到文档: {result['doc_count']} 个")
        
        if result['relevant_docs']:
            print("相关文档:")
            for doc in result['relevant_docs'][:2]:  # 只显示前2个
                print(f"  - {doc['filename']} (相似度: {doc['similarity']:.3f})")
        
        print("-" * 40)
    
    print(f"\n🎉 简易RAG演示完成!")
    print(f"\n💡 优化建议:")
    print("1. 集成真实的LLM API（如阿里云百炼）")
    print("2. 使用更好的向量化模型（如BERT）") 
    print("3. 支持PDF/Word文档解析")
    print("4. 使用向量数据库进行持久化存储")
    print("5. 实现更智能的文档切分策略")

def interactive_rag_chat():
    """交互式RAG对话"""
    
    print("🎯 交互式RAG对话系统")
    print("="*60)
    
    # 初始化系统
    rag_system = SimpleRAGSystem()
    
    # 加载文档
    documents_folder = os.path.join(os.path.dirname(__file__), 'test_documents', 'txt')
    documents = rag_system.load_documents_from_folder(documents_folder)
    
    if not documents:
        print("❌ 没有找到文档，请检查test_documents/txt文件夹")
        return
    
    # 构建索引
    rag_system.build_vector_index(documents)
    
    print(f"\n✅ RAG系统准备就绪！")
    print(f"📄 已加载 {len(rag_system.documents)} 个文档块")
    print(f"💬 开始对话（输入 'quit' 退出）:")
    print("-" * 60)
    
    # 开始交互
    while True:
        try:
            user_input = input("\n🤔 您的问题: ").strip()
            
            if user_input.lower() in ['quit', 'exit', '退出', 'q']:
                print("👋 再见！")
                break
            
            if not user_input:
                continue
            
            # RAG问答
            result = rag_system.chat(user_input, use_llm=True)
            
        except KeyboardInterrupt:
            print("\n\n👋 用户中断，再见！")
            break
        except Exception as e:
            print(f"\n❌ 出现错误: {str(e)}")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='简易RAG系统演示')
    parser.add_argument('--mode', choices=['demo', 'chat'], default='demo',
                       help='运行模式：demo=演示模式，chat=交互模式')
    
    args = parser.parse_args()
    
    if args.mode == 'demo':
        demo_simple_rag()
    elif args.mode == 'chat':
        interactive_rag_chat()
